#!/usr/bin/env python3
# pur_policy_localD.py
# 策略模块：图结构加载、信念管理 + 本地DRL策略计算
import numpy as np
import os
import time
import logging
from math import dist
import networkx as nx

logger = logging.getLogger(__name__)


class PursuitPolicy:
    def __init__(self):
        self.node_coordinates = []  # 节点坐标列表
        self.adj_matrix = []  # 邻接矩阵
        self.dist_3d = []  # 三维距离矩阵
        self.f_matrix = []  # Floyd-Warshall 最短路径矩阵
        self.node_num = 0
        self.degrees = np.zeros(0, dtype=int)  # 节点度数
        self.capture_count = 0
        # 非完美信息参数
        self.obs_dist = 4
        self.obs_nodes = np.array([])
        self.belief_pr = np.zeros(0)  # 初始化为0长度，加载图数据时会重新设置
        self.belief_pos = np.zeros(0, dtype=bool)  # 初始化为0长度，加载图数据时会重新设置
        self.pursuer_mode = 1  # 1: belief preservation
        self.evader_mode = 2  # 2: subsequent move
        # 信念更新帧率限制参数
        self.belief_update_interval = 1.0  # 默认更新间隔为1.0秒
        self.last_belief_update_time = 0.0  # 上次更新时间
        # 加载图数据
        self.load_graph_data()
        logger.info("PursuitPolicy 初始化完成")
        

    def load_graph_data(self):
        """加载图数据（邻接矩阵和预计算的距离矩阵）"""
        root = "../datas/"
        try:
            # 加载邻接矩阵
            adj_matrix_path = os.path.join(root, "adj_matrix.txt")
            if not os.path.exists(adj_matrix_path):
                raise FileNotFoundError(f"邻接矩阵文件不存在: {adj_matrix_path}")
            with open(adj_matrix_path, "r") as f:
                self.node_num = int(f.readline().strip())
                self.adj_matrix = np.zeros((self.node_num, self.node_num), dtype=int)
                self.degrees = np.zeros(self.node_num, dtype=int)
                for i in range(self.node_num):
                    row = list(map(int, f.readline().split()))
                    self.adj_matrix[i] = row
                    self.degrees[i] = np.sum(row)
            # 加载节点坐标
            points_path = os.path.join(root, "points.txt")
            self.load_node_coordinates(points_path)
            # 加载预计算的距离矩阵 - 仅使用NPY格式
            dist_npy_path = os.path.join(root, "D_bin.npy")
            if not os.path.exists(dist_npy_path):
                raise FileNotFoundError(f"距离矩阵文件不存在: {dist_npy_path}")
            # 加载NPY格式
            start_time = time.time()
            self.dist_3d = np.load(dist_npy_path)
            load_time = time.time() - start_time
            # 验证矩阵形状
            if self.dist_3d.shape != (
                self.node_num + 1,
                self.node_num + 1,
                self.node_num + 1,
            ):
                raise ValueError(
                    f"距离矩阵形状不匹配: 预期({self.node_num + 1},{self.node_num + 1},{self.node_num + 1}), "
                    f"实际{self.dist_3d.shape}"
                )
            logger.info(
                f"从NPY文件加载距离矩阵成功! 节点数: {self.node_num}, 耗时: {load_time:.4f}秒"
            )
            # 计算最短路径矩阵 (Floyd-Warshall)
            logger.info("开始计算最短路径矩阵...")
            self.compute_shortest_paths()
            # 初始化信念数组
            self.belief_pr = np.zeros(self.node_num)
            self.belief_pos = np.zeros(self.node_num, dtype=bool)
        except Exception as e:
            logger.error(f"加载图数据错误: {e}")
            raise  # 重新抛出异常以便上层处理

    def load_node_coordinates(self, filename):
        """从文件加载节点坐标，文件不存在时抛出异常"""
        if not os.path.exists(filename):
            raise FileNotFoundError(f"节点坐标文件不存在: {filename}")
        try:
            self.node_coordinates = []
            with open(filename, "r") as f:
                for line in f:
                    # 尝试两种格式：逗号分隔或空格分隔
                    if "," in line:
                        parts = line.strip().split(",")
                    else:
                        parts = line.strip().split()
                    if len(parts) >= 2:
                        x = float(parts[0])
                        y = float(parts[1])
                        self.node_coordinates.append((x, y))
            # 验证节点数量是否匹配
            if len(self.node_coordinates) != self.node_num:
                raise ValueError(
                    f"坐标文件节点数不匹配: 坐标文件 {len(self.node_coordinates)} 个节点, "
                    f"邻接矩阵 {self.node_num} 个节点"
                )
            logger.info(f"成功加载 {len(self.node_coordinates)} 个节点的坐标")
            return True
        except Exception as e:
            logger.error(f"加载节点坐标错误: {e}")
            raise  # 重新抛出异常

    def compute_shortest_paths(self):
        """
        使用 NetworkX 的 floyd_warshall_numpy 加速最短路径矩阵计算
        """
        # 1. 从邻接矩阵构建有向图
        G = nx.from_numpy_array(self.adj_matrix, create_using=nx.DiGraph())
        # 2. 使用 NetworkX 的 floyd_warshall_numpy 计算所有节点对的最短路径
        dist_matrix = nx.floyd_warshall_numpy(G)
        # 3. 定义不可达距离（原代码中为 255）
        INF = 255
        # 4. 替换不可达路径为 INF，并四舍五入为整数
        dist_matrix[np.isinf(dist_matrix)] = INF
        self.f_matrix = np.rint(dist_matrix).astype(int)
        # 5. 确保对角线为 0（自己到自己的距离）
        np.fill_diagonal(self.f_matrix, 0)
        logger.info("使用 NetworkX 成功计算最短路径矩阵")

    def position_to_node(self, x, y):
        """将实际坐标映射到最近的节点"""
        min_dist = float("inf")
        closest_node = 0
        for idx, (nx, ny) in enumerate(self.node_coordinates):
            d = dist((x, y), (nx, ny))
            if d < min_dist:
                min_dist = d
                closest_node = idx
        return closest_node

    def node_to_position(self, node_id):
        """将节点ID转换为实际坐标"""
        if 0 <= node_id < len(self.node_coordinates):
            return self.node_coordinates[node_id]
        return (0.0, 0.0)  # 默认位置

    def get_neighbors(self, node_id):
        """获取节点的邻居（不包括自身）"""
        neighbors = []
        for j in range(self.node_num):
            if self.adj_matrix[node_id][j] == 1:
                neighbors.append(j)
        return neighbors

    def init_belief(self, evader_node):
        """初始化信念状态到evader_node"""
        self.belief_pr.fill(0.0)
        self.belief_pos.fill(False)
        if 0 <= evader_node < self.node_num:
            self.belief_pr[evader_node] = 1.0
            self.belief_pos[evader_node] = True

    def expand_belief(self):
        """传播信念：模拟逃避者随机游走一步（优化版）"""
        node_num = self.node_num
        if np.sum(self.belief_pos) == 0:
            return
        # 优化
        temp = np.zeros(node_num)
        epsilon = 1e-10  # 防零
        stay_prob = 0.5
        move_prob = 0.5

        for i in range(node_num):
            if self.belief_pos[i]:
                pr_i = self.belief_pr[i]
                deg_i = self.degrees[i]
                if deg_i > 0:
                    # 自留
                    temp[i] += pr_i * stay_prob
                    # 移动贡献
                    contrib = pr_i * move_prob / (deg_i + epsilon)
                    for neigh in self.get_neighbors(i):
                        temp[neigh] += contrib
                else:
                    temp[i] += pr_i
        self.belief_pr = np.clip(temp, 0, 1)
        self.belief_pos = self.belief_pr > 1e-8

    def remove_visible(self, p1_node, p2_node):
        """移除从p1或p2可见的位置"""
        for l in range(self.node_num):
            if self.visible(p1_node, p2_node, l):
                self.belief_pos[l] = False
                self.belief_pr[l] = 0.0

    def visible(self, p1, p2, l):
        """检查l是否可见：从p1/p2或obs_nodes距离<=obs_dist"""
        if self.f_matrix[p1][l] <= self.obs_dist:
            return True
        if self.f_matrix[p2][l] <= self.obs_dist:
            return True
        for obs in self.obs_nodes:
            if obs < self.node_num and self.f_matrix[obs][l] <= self.obs_dist:
                return True
        return False

    def normalize_belief(self):
        """归一化信念概率"""
        total = np.sum(self.belief_pr)
        if total > 0:
            self.belief_pr /= total

    def update_belief(self, evader_node, p1_node, p2_node, visible_after_move=False):
        """完整更新信念：传播 -> 移除 -> (if mode) 重置 -> 归一化"""
        # 检查是否需要更新信念（基于时间间隔限制）
        current_time = time.time()
        if current_time - self.last_belief_update_time < self.belief_update_interval:
            # 未达到更新时间间隔，跳过更新
            return  
        self.expand_belief()
        self.remove_visible(p1_node, p2_node)
        if self.pursuer_mode == 0 or visible_after_move:
            self.init_belief(evader_node)
        self.normalize_belief()
        # 更新上次信念更新时间
        self.last_belief_update_time = current_time

    def is_captured(self, evader_node, p1_node, p2_node):
        """检查捕获：基于真实evader_node (服务器知道)，但信念用于策略"""
        # 真实捕获条件
        if evader_node == p1_node or evader_node == p2_node:
            return True
        if (
            self.adj_matrix[evader_node][p1_node] == 1
            or self.adj_matrix[evader_node][p2_node] == 1
        ):
            return True
        return False

    def cond1(self, p1, p2, l):
        """l邻接p1或p2"""
        return self.adj_matrix[l][p1] == 1 or self.adj_matrix[l][p2] == 1

    def _local_calculate_next_move(self, evader_node, p1_node, p2_node):
        """本地计算下一步移动"""
        # 先更新信念 (假设当前信念基于上一步)
        visible_after = self.visible(p1_node, p2_node, evader_node)
        self.update_belief(evader_node, p1_node, p2_node, visible_after)

        # 检查捕获 (真实)
        if self.is_captured(evader_node, p1_node, p2_node):
            return evader_node, p1_node, p2_node

        # 信念捕获检查 (用于策略中断?)
        belief_captured = any(
            self.belief_pos[l] and self.cond1(p1_node, p2_node, l)
            for l in range(self.node_num)
        )
        if belief_captured:
            return evader_node, p1_node, p2_node  # 假设捕获

        evader_neighbors = self.get_neighbors(evader_node)
        p1_neighbors = self.get_neighbors(p1_node)
        p2_neighbors = self.get_neighbors(p2_node)

        # 逃避者决策：后续移动 (mode=2)，基于当前p1,p2，最大化 dist[p1][p2][e_next]
        # (服务器知道真实evader_node，用完全信息模拟逃避者)
        best_evader = evader_node
        max_value = -1
        best_dist_sum = float("-inf")
        for e_next in evader_neighbors:
            value = self.dist_3d[p1_node + 1][p2_node + 1][e_next + 1]
            dist_sum = (
                self.f_matrix[p1_node][e_next]
                + self.f_matrix[p2_node][e_next]
            )
            if value > max_value or (value == max_value and dist_sum > best_dist_sum):
                max_value = value
                best_evader = e_next
                best_dist_sum = dist_sum

        # 追捕者决策：最小化期望最坏距离，基于信念
        best_p1 = p1_node
        best_p2 = p2_node
        min_exp_value = float("inf")
        best_pursuer_dist_sum = float("inf")
        belief_total = np.sum(self.belief_pr[self.belief_pos])
        if belief_total == 0:
            belief_total = 1  # 避免除零

        for p1_next in p1_neighbors:
            for p2_next in p2_neighbors:
                exp_value = 0.0
                for t in range(self.node_num):
                    if self.belief_pos[t]:
                        max_t = 0
                        for tt in self.get_neighbors(t):
                            d = self.dist_3d[p1_next + 1][p2_next + 1][tt + 1]
                            if d > max_t:
                                max_t = d
                        exp_value += self.belief_pr[t] * max_t
                exp_value /= belief_total

                dist_sum = (
                    self.f_matrix[p1_next][evader_node]
                    + self.f_matrix[p2_next][evader_node]
                )
                if exp_value < min_exp_value or (
                    exp_value == min_exp_value and dist_sum < best_pursuer_dist_sum
                ):
                    min_exp_value = exp_value
                    best_p1 = p1_next
                    best_p2 = p2_next
                    best_pursuer_dist_sum = dist_sum

        return best_evader, best_p1, best_p2

    def calculate_next_move(self, evader_node, p1_node, p2_node):
        """计算下一步移动"""
        # 通过本地策略计算下一步 (替换原本地决策逻辑)
        evader_next, pursuer1_next, pursuer2_next = self._local_calculate_next_move(
            evader_node, p1_node, p2_node
        )

        return evader_next, pursuer1_next, pursuer2_next

    def get_belief_pr(self):
        """获取信念概率分布（备份）"""
        return self.belief_pr.copy()

    def get_belief_pos(self):
        """获取信念位置支持（备份）"""
        return self.belief_pos.copy()

    def get_capture_count(self):
        """获取捕获次数"""
        return self.capture_count

    def increment_capture_count(self):
        """增加捕获次数"""
        self.capture_count += 1